{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d469e82c-808e-4577-80f5-c5bc5e9d6107",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-09-13 20:56:11.618412: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-09-13 20:56:13.279428: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "/opt/conda/lib/python3.11/site-packages/xformers/ops/fmha/flash.py:211: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\n",
      "  @torch.library.impl_abstract(\"xformers_flash::flash_fwd\")\n",
      "/opt/conda/lib/python3.11/site-packages/xformers/ops/fmha/flash.py:344: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\n",
      "  @torch.library.impl_abstract(\"xformers_flash::flash_bwd\")\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "from random import choice\n",
    "from diffusers import (\n",
    "    StableDiffusionPipeline,\n",
    "    EulerAncestralDiscreteScheduler,\n",
    "    PNDMScheduler, DPMSolverMultistepScheduler)\n",
    "from tqdm import tqdm\n",
    "\n",
    "from uuid import uuid4\n",
    "from diffusers.utils import load_image, make_image_grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9e8a3758-fb22-4e9a-8cdf-fc93e85b5efa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_images(diffusion_model, syntetic_base_path, steps_number_all=100, resolution=512):\n",
    "    for target, prompt in [\n",
    "        (1, 'a photo of large malignant mole. Malignant'),\n",
    "        (1, 'a photo of large malignant moles. Malignant'),\n",
    "        (1, 'a photo of multiple malignant moles. Malignant'),\n",
    "        (1, 'a photo of multiple lesion moles. Malignant'),\n",
    "        (1, 'a photo of one bening moles and one malignant mole. Malignant')\n",
    "    ]:\n",
    "        for sampler in [PNDMScheduler, DPMSolverMultistepScheduler]:\n",
    "            diffusion_model.scheduler = sampler.from_config(diffusion_model.scheduler.config)\n",
    "    \n",
    "            for i in range(steps_number_all):\n",
    "                image_tmp = diffusion_model(\n",
    "                    prompt=prompt,\n",
    "                    height=resolution,\n",
    "                    width=resolution,\n",
    "                    num_inference_steps=choice([i for i in range(20,40)]),\n",
    "                    guidance_scale = choice([3,4,5,6,7,8]),\n",
    "                    negative_prompt=\"\",\n",
    "                    # generator = [torch.Generator(\"cuda\").manual_seed(1)]\n",
    "                )[0][0]\n",
    "                \n",
    "                overall_index = str(uuid4())\n",
    "                image_tmp.save(os.path.join(syntetic_base_path, f\"{overall_index}___{target}.png\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cad32c24-2b5c-460a-95c0-7e11e1180540",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_folder_to_save = \"../data/artifacts/syntetic_base_folds_final\"\n",
    "os.makedirs(base_folder_to_save, exist_ok=True)\n",
    "\n",
    "finetunes_path = \"...\" #path to your finetunes_dir\n",
    "finetunes_epochs_to_use = ['epoch-000035', 'epoch-000040', 'epoch-000045', 'epoch-000050']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ad0dfbf-e668-49a6-9f73-065f9ca1a748",
   "metadata": {},
   "outputs": [],
   "source": [
    "images_list = []\n",
    "for finetune_name in os.listdir(finetunes_path):\n",
    "    tmp = [i for i in finetune_name.split('_') if '' != i]\n",
    "    seed_tmp = tmp[-3] #-2\n",
    "    fold_tmp = tmp[-1]\n",
    "    print(seed_tmp, fold_tmp)\n",
    "\n",
    "    base_folder_to_save_tmp_hr = os.path.join(base_folder_to_save, seed_tmp, fold_tmp, \"hr\")\n",
    "    os.makedirs(base_folder_to_save_tmp_hr, exist_ok=True)\n",
    "\n",
    "    base_folder_to_save_tmp_lr = os.path.join(base_folder_to_save, seed_tmp, fold_tmp, \"lr\")\n",
    "    os.makedirs(base_folder_to_save_tmp_lr, exist_ok=True)\n",
    "    for epoch_model in finetunes_epochs_to_use:\n",
    "        diffusion_model = StableDiffusionPipeline.from_pretrained( #from_single_file\n",
    "            # \"Derm-T2IM.safetensors\",\n",
    "            os.path.join(finetunes_path, finetune_name, epoch_model),\n",
    "            torch_dtype=torch.float16,\n",
    "            safety_checker=None,\n",
    "            use_safetensors=True\n",
    "        )\n",
    "        diffusion_model = diffusion_model.to(\"cuda:0\")\n",
    "        \n",
    "        generate_images(diffusion_model, base_folder_to_save_tmp_hr, steps_number_all=100, resolution=512)\n",
    "        \n",
    "        generate_images(diffusion_model, base_folder_to_save_tmp_lr, steps_number_all=100, resolution=128)\n",
    "        \n",
    "        diffusion_model.to('cpu')\n",
    "        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
